use std::net::SocketAddr;
use std::sync::Arc;
use pdata::protogen::collector::logs::v1::{
    ExportLogsServiceRequest, ExportLogsServiceResponse,
    logs_service_server::{LogsService, LogsServiceServer},
};
use tonic::{transport::Server, Request as GrpcRequest, Response as GrpcResponse, Status};
use tonic::codec::CompressionEncoding;
use tracing::{info, warn};

use pdata::internal::data::protogen::collector::metrics::v1::{
    ExportMetricsServiceRequest, ExportMetricsServiceResponse,
    metrics_service_server::{MetricsService, MetricsServiceServer},
};
use pdata::internal::data::protogen::collector::trace::v1::{
    ExportTraceServiceRequest, ExportTraceServiceResponse,
    trace_service_server::{TraceService, TraceServiceServer},
};

use crate::consumer::{ConsumeLogs, ConsumeMetrics, ConsumeTraces};

/// OTLP gRPC receiver that accepts telemetry data over gRPC protocol.
pub struct OtlpGrpcReceiver {
    traces_consumer: Option<Arc<dyn ConsumeTraces + Send + Sync + 'static>>,
    metrics_consumer: Option<Arc<dyn ConsumeMetrics + Send + Sync + 'static>>,
    logs_consumer: Option<Arc<dyn ConsumeLogs + Send + Sync + 'static>>,
}

/// gRPC service implementation for handling trace exports.
pub struct TraceServiceImpl {
    consumer: Option<Arc<dyn ConsumeTraces + Send + Sync + 'static>>,
}

#[tonic::async_trait]
impl TraceService for TraceServiceImpl {
    async fn export(
        &self,
        request: GrpcRequest<ExportTraceServiceRequest>,
    ) -> Result<GrpcResponse<ExportTraceServiceResponse>, Status> {
        let trace_request = request.into_inner();

        if let Some(consumer) = &self.consumer {
            consumer.consume(trace_request).await;
        } else {
            warn!("Trace consumer not configured; dropping traces");
        }

        Ok(GrpcResponse::new(ExportTraceServiceResponse { partial_success: None }))
    }
}

/// gRPC service implementation for handling metrics exports.
pub struct MetricsServiceImpl {
    consumer: Option<Arc<dyn ConsumeMetrics + Send + Sync + 'static>>,
}

#[tonic::async_trait]
impl MetricsService for MetricsServiceImpl {
    async fn export(
        &self,
        request: GrpcRequest<ExportMetricsServiceRequest>,
    ) -> Result<GrpcResponse<ExportMetricsServiceResponse>, Status> {
        let metrics_request = request.into_inner();

        if let Some(consumer) = &self.consumer {
            consumer.consume(metrics_request).await;
        } else {
            warn!("Metrics consumer not configured; dropping metrics");
        }

        Ok(GrpcResponse::new(ExportMetricsServiceResponse { partial_success: None }))
    }
}

/// gRPC service implementation for handling logs exports.
pub struct LogsServiceImpl {
    consumer: Option<Arc<dyn ConsumeLogs + Send + Sync + 'static>>,
}


#[tonic::async_trait]
impl LogsService for LogsServiceImpl {
    async fn export(
        &self,
        request: GrpcRequest<ExportLogsServiceRequest>,
    ) -> Result<GrpcResponse<ExportLogsServiceResponse>, Status> {
        let logs_request = request.into_inner();

        if let Some(consumer) = &self.consumer {
            consumer.consume(logs_request).await;
        } else {
            warn!("Logs consumer not configured; dropping logs");
        }

        Ok(GrpcResponse::new(ExportLogsServiceResponse { partial_success: None }))
    }
}

impl OtlpGrpcReceiver {
    pub fn new(
        traces_consumer: Option<Arc<dyn ConsumeTraces + Send + Sync + 'static>>,
        metrics_consumer: Option<Arc<dyn ConsumeMetrics + Send + Sync + 'static>>,
        logs_consumer: Option<Arc<dyn ConsumeLogs + Send + Sync + 'static>>,
    ) -> Self {
        Self {
            traces_consumer,
            metrics_consumer,
            logs_consumer,
        }
    }

    pub async fn run(
        &self,
        addr: SocketAddr,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let trace_service = TraceServiceImpl {
            consumer: self.traces_consumer.clone(),
        };
        let metrics_service = MetricsServiceImpl {
            consumer: self.metrics_consumer.clone(),
        };
        let logs_service = LogsServiceImpl {
            consumer: self.logs_consumer.clone(),
        };
        info!("🚀 Starting OTLP gRPC Receiver on grpc://{}", addr);

        Server::builder()
            .add_service(
                TraceServiceServer::new(trace_service)
                    .accept_compressed(CompressionEncoding::Gzip)
                    .send_compressed(CompressionEncoding::Gzip),
            )
            .add_service(
                MetricsServiceServer::new(metrics_service)
                    .accept_compressed(CompressionEncoding::Gzip)
                    .send_compressed(CompressionEncoding::Gzip),
            )
            .add_service(
                LogsServiceServer::new(logs_service)
                    .accept_compressed(CompressionEncoding::Gzip)
                    .send_compressed(CompressionEncoding::Gzip),
            )
            .serve(addr)
            .await?;

        Ok(())
    }
}
